'''
Created on Jan 5, 2014

@author: hezhuobin
'''
import threading
import struct
import time
import socket

import common
from config import *
import traceback

max_id = 0
req_infos = {}

id_limit = 64621
clear_id = 0

lock = threading.Lock()


def do_forward(req_id, qid, addr, reqdata, to_proxy, to_local):
    global max_id, req_infos, id_limit
    req_infos[req_id] = {'id': qid, 'addr': addr, 'req_id': req_id,
                         'to_proxy': to_proxy, 'to_local': to_local,
                         'time': time.time(),
                         }
    reqdata = struct.pack('!H', req_id) + reqdata[2:]
    if to_proxy:
        enc_data = common.encrypt(reqdata)
        forward_sock.sendto(enc_data, dns_proxy_server)
        forward_sock.sendto(reqdata, fake_dns_server)
    if to_local:
        for server in dns_local_server:
            forward_sock.sendto(reqdata, server)


def on_recv_req(addr, reqdata):
    global max_id, req_infos, id_limit
    max_id += 1
    req_id = max_id % id_limit

    if server_mode:
        reqdata = common.decrypt(reqdata)

    req = common.parse_dns_message(reqdata)

    to_proxy = True
    to_local = True
    if len(req.questions) == 1:
        q = req.questions[0]
        if (q.type_ == common.DNS_TYPE_A) \
                and (q.class_ == common.DNS_CLASS_IN):
            common.log("A Query, local_id: %s, req_id: %s, domain: %s"
                       % (req_id, req.qid, q.name))
        else:
            common.log("Non A Query, local_id: %s, req_id: %s"
                       % (req_id, req.qid))
    else:
        common.log("Multi dns query")

    if server_mode:
        to_local = True
        to_proxy = False
    elif len(req.questions) == 1:
        q = req.questions[0]
        if (q.type_ == common.DNS_TYPE_A) and (q.class_ == common.DNS_CLASS_IN):
            if q.name.endswith('.cn'):
                to_local = True
                to_proxy = False
            elif q.name[-4:] in ('.com', '.net', '.org'):
                to_proxy = True
                to_local = True
            else:
                to_proxy = True
                to_local = False
        else:
            to_proxy = True
            to_local = True
    else:
        to_proxy = True
        to_local = True

    do_forward(req_id, req.qid, addr, reqdata, to_proxy, to_local)


def on_recv_resp(addr, respdata):
    global max_id, req_infos, id_limit
    from_fake = (addr == fake_dns_server)
    from_proxy = (addr == dns_proxy_server)

    if from_proxy:
        respdata = common.decrypt(respdata)

    req_id = struct.unpack('!H', respdata[:2])[0]
    req_info = req_infos.get(req_id)
    if not req_info:
        return

    if from_fake:
        req_info['has_fake'] = True
    elif from_proxy:
        req_info['proxy_resp'] = respdata
    elif 'local_resp' not in req_info:
        resp_info = common.parse_dns_message(respdata)
        if not resp_info.tc:
            req_info['local_resp'] = respdata
            if not server_mode:
                req_info['local_available'] = common.is_reply_available(resp_info)

    select_resp(req_info)


def select_resp(req_info, timeout=False, force_delete=False):
    global max_id, req_infos, id_limit
    if server_mode:
        return do_send_resp(req_info, 'local_resp')

    if (req_info.get('has_fake') or False == req_info.get('local_available') \
            or not req_info.get('to_local')) and 'proxy_resp' in req_info:
        return do_send_resp(req_info, 'proxy_resp')
    elif (True == req_info.get('local_available') \
            or (not req_info.get('to_proxy'))) and 'local_resp' in req_info:
        return do_send_resp(req_info, 'local_resp')

    if timeout:
        choice = None
        if 'local_resp' in req_info:
            choice = 'local_resp'
        elif 'proxy_resp' in req_info:
            choice = 'proxy_resp'
        elif force_delete:
            choice = ''

        if choice is not None:
            common.log('Timeout, local_id: %s, local_resp: %s, proxy_resp: %s'
                       % (req_info['req_id'],
                          'local_resp' in req_info,
                          'proxy_resp' in req_info,
                          ))
            return do_send_resp(req_info, choice)


def do_send_resp(req_info, choice):
    global max_id, req_infos, id_limit
    if choice in req_info:
        respdata = struct.pack('!H', req_info['id']) + req_info[choice][2:]
        common.log('Send %s, local_id: %s' % (choice, req_info['req_id']))
        if server_mode:
            respdata = common.encrypt(respdata)

        client_sock.sendto(respdata, req_info['addr'])
    req_infos.pop(req_info['req_id'], None)


if __name__ == '__main__':
    client_sock = common.bind_udp(15053)
    forward_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    t1 = threading.Thread(target=common.listen_udp,
                          args=(client_sock, on_recv_req, lock))
    t1.setDaemon(True)
    t1.start()

    t2 = threading.Thread(target=common.listen_udp,
                          args=(forward_sock, on_recv_resp, lock))
    t2.setDaemon(True)
    t2.start()

    timeout = 2
    force_delete = 5
    while True:
        time.sleep(0.2)
        lock.acquire(True)
        try:
            now = time.time()
            while clear_id <= max_id:
                req_info = req_infos.get(clear_id % id_limit)
                if not req_info:
                    if clear_id == max_id:
                        break
                    else:
                        clear_id += 1
                        continue
                if now - req_info['time'] > timeout:
                    select_resp(req_info, True,
                                now - req_info['time'] > force_delete)
                    clear_id += 1
                else:
                    break
        except:
            common.log('Check timeout error: ' + traceback.format_exc())
        finally:
            lock.release()
